import copy
import socket
import threading
import sys
import os
import struct

import numpy
import torch
from torch import device, nn

from LeNet_model_MNIST import LeNet_edge_side, LeNet_client_side
from get_data import get_mnist, dataset_iid

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
epochs = 2
lr = 0.001

if_train = 0  # 璁粌鏍囧織
connected = 0  # 杩炴帴缁堢涓暟
model = 0  # 鎺ユ敹灞€閮ㄦā鍨嬩釜鏁�
model_kind = '0'  # 妯″瀷绉嶇被
lock_con = threading.Lock()
lock_model = threading.Lock()
lock_kind = threading.Lock()
lock_accuracy = threading.Lock()
lock_epoch = threading.Lock()
client_model = 0  # 缁堢妯″瀷绉嶇被鍚屾
net = 0
is_agg = 0
Net = {}
epoch = 4
final_client_path = '0'
final_server_path = '0'
final_model = 0
accuracy = []
accuracy_epoch = {}
epoch_over = 0

criterion = nn.CrossEntropyLoss()


def calculate_accuracy(fx, y):
    preds = fx.max(1, keepdim=True)[1]
    correct = preds.eq(y.view_as(preds)).sum()
    acc = 100.00 * correct.float() / preds.shape[0]
    return acc


def send_string(con, string):  # 鍚戝鎴风鍙戦€佸叧闂寚浠�
    con.send(string.encode())


def send_file(con, file_addr):
    while True:
        filepath = file_addr  # 浼犺緭鏂囦欢鐨勮矾寰�
        if os.path.isfile(filepath):
            # 瀹氫箟瀹氫箟鏂囦欢淇℃伅銆�128s琛ㄧず鏂囦欢鍚嶄负128bytes闀匡紝l琛ㄧず涓€涓猧nt鎴杔og鏂囦欢绫诲瀷锛屽湪姝や负鏂囦欢澶у皬
            fileinfo_size = struct.calcsize('128sq')
            # 瀹氫箟鏂囦欢澶翠俊鎭紝鍖呭惈鏂囦欢鍚嶅拰鏂囦欢澶у皬
            fhead = struct.pack('128sq', bytes(os.path.basename(filepath).encode('utf-8')), os.stat(filepath).st_size)
            con.send(fhead)
            print('client filepath: {0}'.format(filepath))
            with open(filepath, 'rb') as fp:
                while True:
                    data = fp.read(1024)
                    if not data:
                        print('{0} file send over...'.format(filepath))
                        break
                    con.send(data)
        break


def deal_data(conn, port):
    print('Accept new connection from {0}'.format(port))
    while True:
        fileinfo_size = struct.calcsize('128sq')  # linux 鍜� windows 浜掍紶 128sl 鏀逛负 128sq  鏈哄櫒浣嶆暟涓嶄竴鏍凤紝涓€涓�32浣嶄竴涓�64浣�
        buf = conn.recv(fileinfo_size)
        print('鏀跺埌鐨勫瓧鑺傛祦锛�', buf, type(buf))
        if buf:
            print(buf, type(buf))
            filename, filesize = struct.unpack('128sq', buf)
            fn = filename.strip(str.encode('\00'))
            new_filename = os.path.join(str.encode('./'), str.encode(port) + fn)  # 鏀跺埌鐨勬枃浠跺懡鍚嶏紝鎸夌収闇€姹傛洿鏀�
            print('file new name is {0}, filesize if {1}'.format(new_filename, filesize))
            recvd_size = 0  # 瀹氫箟宸叉帴鏀舵枃浠剁殑澶у皬
            with open(new_filename, 'wb') as fp:
                print("start receiving...")
                while not recvd_size == filesize:
                    if filesize - recvd_size > 1024:
                        data = conn.recv(1024)
                        recvd_size += len(data)
                    else:
                        data = conn.recv(filesize - recvd_size)
                        recvd_size = filesize

                    fp.write(data)
            print("end receive...")
            return new_filename.decode()
        break


def socket_service(port):
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        # 淇敼ip,姝ゅip蹇呴』涓烘湇鍔″櫒绔殑ip ,linux鍋氭湇鍔″櫒杈撳叆ifconfig寰楀埌ip
        s.bind(('0.0.0.0', port))
        s.listen(10)

    except socket.error as msg:
        print(msg)
        sys.exit(1)
    print("Waiting...")
    while True:
        conn, addr = s.accept()
        t = threading.Thread(target=server_client, args=(conn, str(port)))
        t.start()
        break


def server_cloud(ip_addr, port):
    try:
        conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # 淇敼ip
        conn.connect((ip_addr, port))  # 姝ゅip蹇呴』涓烘湇鍔″櫒绔殑ip

    except socket.error as msg:
        print(msg)
        sys.exit(1)

    while True:
        if connected == 2:
            send_string(conn, 'ready')
            break

    global model_kind, net, if_train, final_client_path, final_server_path, final_model

    while True:
        if client_model == 2:
            send_string(conn, 'all-ready' + str(port))
            print('all-ready')
            break

    command = conn.recv(1024).decode()

    if command == 'train':
        if_train = 1

    a = []
    while True:
        if epoch_over == 2:
            for key in accuracy_epoch.keys():
                a.append(accuracy_epoch[key])
            b = numpy.array(a)
            b = numpy.mean(b, axis=0)
            for value in b:
                send_string(conn, str(value))
                while True:
                    command = conn.recv(1024).decode()
                    if command == 'recvd':
                        break
            break
    send_string(conn, 'over')

    while True:
        if os.path.exists('./client.pth') and os.path.exists('./server.pth'):
            send_file(conn, './client.pth')
            while True:
                command = conn.recv(1024).decode()
                if command == 'recvd':
                    break
            send_file(conn, './server.pth')
            break

    final_client_path = deal_data(conn, port)
    send_string(conn, 'recvd')
    final_server_path = deal_data(conn, port)

    final_model = 1

    while True:
        if len(accuracy) == 2:
            for a in accuracy:
                send_string(conn, a)
                while True:
                    command = conn.recv(1024).decode()
                    if command == 'recvd':
                        break
            break

    send_string(conn, 'over')
    conn.close()


def server_client(conn, port):
    global connected, accuracy, accuracy_epoch, epoch_over
    lock_con.acquire()
    connected += 1
    lock_con.release()

    command = conn.recv(1024).decode()
    global client_model
    if command == 'ready':
        lock_kind.acquire()
        client_model += 1
        lock_kind.release()

    while True:
        if if_train == 1:
            break

    net_glob_edge = LeNet_edge_side()
    if torch.cuda.device_count() > 1:
        print("We use", torch.cuda.device_count(), "GPUs")
        net_glob_edge = nn.DataParallel(net_glob_edge)
    accuracy_epoch[port] = []

    net_glob_edge.to(device)
    print(net_glob_edge)
    for i in range(epoch):
        send_string(conn, 'train')
        while True:
            command = conn.recv(1024).decode()
            print(command)
            if command == 'finish':
                data = deal_data(conn, port)
                send_string(conn, 'recvd')
                label = deal_data(conn, port)
                inputs = torch.load(data)
                label = torch.load(label)

                net_glob_edge.train()
                optimizer_edge = torch.optim.Adam(net_glob_edge.parameters(), lr)

                optimizer_edge.zero_grad()

                inputs = inputs.to(device)
                label = label.to(device)
                inputs.required_grad = True
                inputs.retain_grad()
                # ---------forward prop-------------
                fx_edge = net_glob_edge(inputs)
                print('fx_edge:', fx_edge)
                # calculate loss
                loss = criterion(fx_edge, label)
                print('loss:', loss)
                # calculate accuracy
                acc = calculate_accuracy(fx_edge, label)
                print(acc)
                # --------backward prop--------------
                loss.backward()
                dfx_client = inputs.grad.clone().detach()
                optimizer_edge.step()

                dfx = dfx_client
                print('dfx:', dfx)

                accuracy_epoch[port].append(acc)

                torch.save(dfx, './grad.pt')
                send_file(conn, './grad.pt')
            elif command == 'over':
                break
        if i == 1:
            deal_data(conn, port)
            torch.save(net_glob_edge.state_dict(), './' + port + 'phase_server.pth')
            while True:
                if os.path.exists('./phase_client.pth') and os.path.exists('./phase_server.pth'):
                    send_file(conn, './phase_client.pth')
                    net_glob_edge.load_state_dict(torch.load('./phase_server.pth'))
                    break
    send_string(conn, 'over')
    lock_epoch.acquire()
    epoch_over += 1
    lock_epoch.release()
    torch.save(net_glob_edge.state_dict(), './' + port + 'server.pth')
    deal_data(conn, port)

    while True:
        if final_model == 1:
            send_file(conn, final_client_path)
            break

    final_server_model = LeNet_edge_side()
    final_server_model.load_state_dict(torch.load(final_server_path))

    with torch.no_grad():
        data = deal_data(conn, port)
        send_string(conn, 'recvd')
        label = deal_data(conn, port)
        test_inputs = torch.load(data)
        test_label = torch.load(label)
        output = final_server_model(test_inputs)
        accuracy_ff = calculate_accuracy(test_label, output)

    lock_accuracy.acquire()
    accuracy.append(str(accuracy_ff))
    lock_accuracy.release()


def FedAvg(w):
    w_avg = copy.deepcopy(w[0])
    for p in w_avg.keys():
        for i in range(1, len(w)):
            w_avg[p] += w[i][p]
        w_avg[p] = torch.div(w_avg[p], len(w))
    return w_avg


if __name__ == '__main__':
    t1 = threading.Thread(target=server_cloud, args=('10.251.139.40', 5558))
    t1.start()
    #server_cloud('10.251.139.40', 5558)
    t2 = threading.Thread(target=socket_service, args=[5555])
    t2.start()
    t3 = threading.Thread(target=socket_service, args=[5556])
    t3.start()
    #socket_service(5555)
    #socket_service(5556)


    while True:
        if os.path.exists('./5555phase_client.pth') and os.path.exists('./5556phase_client.pth'):
            '''
            鑱氬悎绠楁硶
            鍋囪缁撴灉涓篶lient
            '''
            net1 = LeNet_client_side()
            net2 = LeNet_client_side()
            net1.load_state_dict(torch.load('./5555phase_client.pth'))
            net2.load_state_dict(torch.load('./5556phase_client.pth'))
            w = [net1.state_dict(), net2.state_dict()]

            phase_model = FedAvg(w)

            torch.save(phase_model, './phase_client.pth')
            os.remove('./5555phase_client.pth')
            os.remove('./5556phase_client.pth')
            break

    while True:
        if os.path.exists('./5555phase_server.pth') and os.path.exists('./5556phase_server.pth'):
            '''
            鑱氬悎绠楁硶
            鍋囪缁撴灉涓簊erver
            '''
            net1 = LeNet_edge_side()
            net2 = LeNet_edge_side()
            net1.load_state_dict(torch.load('./5555phase_server.pth'))
            net2.load_state_dict(torch.load('./5556phase_server.pth'))
            w = [net1.state_dict(), net2.state_dict()]

            phase_model = FedAvg(w)

            torch.save(phase_model, './phase_server.pth')
            os.remove('./5555phase_server.pth')
            os.remove('./5556phase_server.pth')
            break

    while True:
        if os.path.exists('./5555client.pth') and os.path.exists('./5556client.pth'):
            '''
            鑱氬悎绠楁硶
            鍋囪缁撴灉涓篶lient
            '''
            net1 = LeNet_client_side()
            net2 = LeNet_client_side()
            net1.load_state_dict(torch.load('./5555client.pth'))
            net2.load_state_dict(torch.load('./5556client.pth'))
            w = [net1.state_dict(), net2.state_dict()]

            final = FedAvg(w)
            torch.save(final, './client.pth')
            break

    while True:
        if os.path.exists('./5555server.pth') and os.path.exists('./5556server.pth'):
            '''
            鑱氬悎绠楁硶
            鍋囪缁撴灉涓簊erver
            '''
            net1 = LeNet_edge_side()
            net2 = LeNet_edge_side()
            net1.load_state_dict(torch.load('./5555server.pth'))
            net2.load_state_dict(torch.load('./5556server.pth'))
            w = [net1.state_dict(), net2.state_dict()]

            phase_model = FedAvg(w)

            torch.save(phase_model, './server.pth')
            break